package systemf;

trait DeBruijnSubstitution extends SubstitutionProvider {
  import DeBruijnOps._

  
  type VarID=String
  type VarRef=Int
  
  type partial = bruijnpart
  //type full = bruijnfull
  type topsubst = bruijntopsubst
  override def substituteterm(v:Term)=new bruijnpart(v)
  override def substitutetype(v:Type)=new bruijnpart(v)
  
  override def removingBinder(ty:Type):Type = {
    val typ = typeShift(-1,ty)
    val unit = typeMap({case (int,TyVar(n)) => if (n <0) throw new IllegalArgumentException("scope error") else TyUnit()}, 0, typ)
    typ
  }
  def addingBinder(ty:Type):Type = {
    typeShift(+1, ty)
  }
  
  //def underBinder(f: =>Type):Type = typeShift(-1, f)
  
  
  class bruijnpart(v:Term) extends examplepartial(v) {
    override def asTop = new bruijntopsubst(v)
    
  }
   class bruijntopsubst(v:Term) extends exampletopsubst(v) {
    override def intoterm(term: Term)= (v,term) match {
      //case (v:Type,ty:Type) => typeintotypeSubstTop(v,ty)
      case (v:Type,t:Term)  => typeintotermSubstTop(v,t) // tell compiler that since term matches with LTerm, that a return type of LTerm is ok
      case (v:Term,t:Term)  => termSubstTop(v,t)
    }
    override def intotype(ty: Type) = (v,ty) match {
      case (v:Type, ty:Type) => typeintotypeSubstTop(v,ty)
      case (v:Term, ty:Type) => throw new Error("cannot form dependent type")
    }
  }
    

}




object DeBruijnOps {
  type Term = LTerm
  type Type = LType
  
  def termSubstTop(v: Term, t: Term): Term = {
    termShift(-1, 
              {val termsubst= termSubst(0, 
                        {val valunderbinder = termShift(1,v);valunderbinder},
                        t );termsubst} )
  }
  def typeintotermSubstTop(v: Type, t: Term) : Term= {
    termShift(-1,
               typeintotermSubst(0,
                         typeShift(1,v),
                         t) )
  }
  def typeintotypeSubstTop(v: Type, t:Type) : Type = {
    typeShift(-1,
               {val typesubst = typeSubst(0,
                       {val tyunderbinder = typeShift(1,v)
                             ; tyunderbinder},
                       t) ;typesubst} )
  }
  
  def termShift(d: Int, t: Term) = {
    termMap((c:Int,x:Var) => if (x.n >=c) Var(x.n+d) else Var(x.n) ,
            (c:Int,x:Type) => typeShiftAbove(d,c,x) ,
            0,t) //initial cutoff is 0
  }
  
  
  def typeShiftAbove(d:Int, c:Int, ty : Type) :Type = {
    typeMap((c:Int,x:TyVar) => if (x.n >= c) TyVar(x.n +d) else TyVar(x.n), c, ty)
  }
  
  def typeShift(d:Int,ty:Type):Type = 
    typeMap((c:Int,x:TyVar) => if (x.n >= c) TyVar(x.n +d) else TyVar(x.n), 0, ty)

  
  def termSubst(j: Int, s: Term,t: Term) = {
    termMap((c:Int,x:Var) => if (x.n == j+c) termShift(c,s) else Var(x.n),
            (c:Int,x:Type) => x,
            j,t)
  }
  def typeSubst(j:Int, s: Type, ty: Type) = {
    typeMap((c:Int,x:TyVar) => if (x.n == c) 
                                   typeShift(c,s) 
                               else TyVar(x.n),j,ty)
  }
  
  
  def typeintotermSubst(j:Int, s:Type, t:Term) = {
    termMap( (c:Int,x:Var) => x,
             (c:Int,x:Type) => typeSubst(c,s,x), //CAREFUL: thread c, not j: c is updated to be in sync with the depth of the binding
             j,t)
  }
  
  
  def termMap(onvar:(Int,Var)=>Term,ontype:(Int,Type)=>Type, i:Int, t:Term) ={
    //TODO: can you generalise this into a library for dealing with de Bruijn indices?
    // i.e., clients of that library should just write a map that calls the right function for every AST construct
    // (--> provide helper for binders and variables)
       
    def walk(c: Int, t: Term) : Term = t match {
      case Var(nb)        => onvar(c,Var(nb))
      case Lam(x,ty,t1)   => Lam(x,ontype(c,ty),walk(c+1,t1)) //+1 because we are looking inside a new scope
      case App(t1,t2)     => App(walk(c,t1),walk(c,t2))
      
      case Let(x,t1,t2) => Let(x,walk(c,t1), walk(c+1, t2))
      //case TyVar(nb)    => ontype(c,TyVar(nb))
      case TLam(x,t)      => TLam(x,walk(c+1,t))
      case TApp(t,ty)     => TApp(walk(c,t), ontype(c, ty))
      case Pack(ty, t, pubty) => Pack(ontype(c,ty),walk(c,t),ontype(c,pubty))
      case Unpack(impl, pubtyn,pubvarn, body) =>
                           Unpack(walk(c,impl),pubtyn,pubvarn,walk(c+2,body)) //+2 because we have introduced two variables here
    
      case Record(fields)=> {
        val fields_ = fields.map({case (label,fld) => (label, walk(c,fld))})
        Record(fields_)
      }
      case RecProj(r,label)  => RecProj(walk(c,r),label)
      
      case base:BaseValue=> base
      case If(cond,p,n)  => If(walk(c,cond),walk(c,p),walk(c,n))  
    }
    walk(0,t)
  }
  
  def typeMap(ontyvar:(Int,TyVar)=>Type, c:Int, ty: Type) = {
    def walk(c:Int, ty:Type):Type = ty match {
    case TyVar(nb) => ontyvar(c,TyVar(nb)) //if (nb >= c) TyVar(nb+d) else TyVar(nb) 
    
    case ty:BaseType => ty

    case TyArr(t1, t2) =>
      val t1_ = walk(c,t1)
      val t2_ = walk(c,t2)
     TyArr(t1_, t2_)
    case TyUniv(tyX, ty) =>
     TyUniv(tyX, walk(c+1,ty))  //extra var bound, so shift cutoff by one or we will take the newly introduced one
    case TyExist(tyX, ty) =>
     TyExist(tyX, walk(c+1,ty))
    case TyRecord(fields)=> {
      val fields_ = fields.map({case (label,fld)=> (label, walk(c,fld))})
      TyRecord(fields_)
    }   
    }
    
    walk(c,ty)
  }
}
